/**
 * \file lm_ellipse.cpp
 * \author Zhihao Zhan (zhanzhihao_dt@163.com)
 * \brief lm拟合椭圆
 * \version 0.1
 * \date 2023-03-21
 *
 * @copyright Copyright (c) 2023
 *
 */

#include <iostream>
#include <fstream>
#include <vector>
#include <iomanip>
#include <cmath>

#include <eigen3/Eigen/Core>
#include <eigen3/Eigen/Dense>

#include <opencv2/opencv.hpp>

#define max(a, b) (((a) > (b)) ? (a) : (b))

Eigen::MatrixXd fx_; // residual
Eigen::MatrixXd J_;  // jacbian

double Func(const Eigen::VectorXd &obj)
{
  return obj.squaredNorm() / 2;
}

double CostFunc(const std::vector<cv::Point2d> &input, const Eigen::VectorXd &params, int index)
{
  double      a  = params(0);
  double      b  = params(1);
  double      c  = params(2);
  double      d  = params(3);
  double      e  = params(4);
  double      f  = params(5);
  cv::Point2d pt = input.at(index);

  return a * pt.x * pt.x + b * pt.y * pt.y + c * pt.x * pt.y + d * pt.x + e * pt.y + f;
}

void ObjF(const std::vector<cv::Point2d> &input, const Eigen::VectorXd &params)
{
  fx_.resize(input.size(), 1);
  for (int i = 0; i < input.size(); ++i)
  {
    fx_(i) = CostFunc(input, params, i);
  }
}

void Jacobian(const std::vector<cv::Point2d> &input, const Eigen::VectorXd &params)
{
  int rowNum = input.size();
  int colNum = params.rows();
  J_.resize(rowNum, colNum);
  double a = params(0);
  double b = params(1);
  double c = params(2);
  double d = params(3);
  double e = params(4);
  double f = params(5);

  for (int i = 0; i < rowNum; i++)
  {
    cv::Point pt = input.at(i);

    J_(i, 0) = pt.x * pt.x;
    J_(i, 1) = pt.y * pt.y;
    J_(i, 2) = pt.x * pt.y;
    J_(i, 3) = pt.x;
    J_(i, 4) = pt.y;
    J_(i, 5) = 1;
  }
}

double MaxMatrixDiag(const Eigen::MatrixXd &Hessian)
{
  int max = 0;
  for (int i = 0; i < Hessian.rows(); i++)
  {
    if (Hessian(i, i) > max)
      max = Hessian(i, i);
  }

  return max;
}

double DeltaL(const Eigen::VectorXd &h)
{
  Eigen::MatrixXd L = -h.transpose() * J_.transpose() * fx_ - 0.5 * h.transpose() * J_.transpose() * J_ * h;
  return L(0, 0);
}

void LM(const std::vector<cv::Point2d> &input, int iter_times)
{
  // step: 1. 初始化参数
  int             param_nums = 6;
  Eigen::VectorXd params_lm(param_nums);
  // init gauss
  params_lm << 3.5, 5, 0, 1.7, 12.0, 2;

  ObjF(input, params_lm);
  Jacobian(input, params_lm);

  Eigen::MatrixXd Hessian  = J_.transpose() * J_;  // Hessian
  Eigen::VectorXd gradient = J_.transpose() * fx_; // gradient
  // initial parameter tao v epsilon1 epsilon2
  double    tao  = 1e-3;
  long long v    = 2;
  double    eps1 = 1e-12, eps2 = 1e-12;
  double    u     = tao * MaxMatrixDiag(Hessian);
  bool      found = gradient.norm() <= eps1;
  if (found)
  {
    std::cout << "Levenberg-Marquardt parameter: " << std::endl << params_lm << std::endl << std::endl << std::endl;
    return;
  }

  double last_sum = 0;
  int    iterCnt  = 0;

  // step: 2. 迭代优化
  while (iterCnt < iter_times)
  {
    ObjF(input, params_lm);
    Jacobian(input, params_lm);      // jacobin
    Hessian  = J_.transpose() * J_;  // Hessian
    gradient = J_.transpose() * fx_; // gradient

    if (gradient.norm() <= eps1)
    {
      std::cout << "stop g(x) = 0 for a local minimizer optimizer." << std::endl;
      break;
    }

    // std::cout << "Hessian: " << std::endl << Hessian << std::endl;
    // std::cout << "Gradient: " << std::endl << gradient << std::endl;

    Eigen::VectorXd step =
        (Hessian + u * Eigen::MatrixXd::Identity(param_nums, param_nums)).inverse() * gradient; // negtive Hlm.

    std::cout << "step: " << std::endl << step << std::endl;

    if (step.norm() <= eps2 * (params_lm.norm() + eps2))
    {
      std::cout << "stop because change in x is small" << std::endl;
      break;
    }

    Eigen::VectorXd paramsNew(params_lm.rows());
    paramsNew = params_lm - step; // h_lm = -step;

    // compute f(x)
    ObjF(input, params_lm);
    Eigen::VectorXd fx(input.size(), 1);
    fx = fx_;

    // compute f(x_new)
    Eigen::VectorXd fx_new(input.size(), 1);
    ObjF(input, paramsNew);
    fx_new = fx_;

    double deltaF = Func(fx) - Func(fx_new);
    double deltaL = DeltaL(-1 * step);

    double roi = deltaF / deltaL;
    std::cout << "roi is : " << roi << std::endl;
    if (roi > 0)
    {
      params_lm = paramsNew;
      u *= max(1.0 / 3.0, 1 - pow(2 * roi - 1, 3));
      v = 2;
    } else
    {
      u = u * v;
      v = v * 2;
    }

    std::cout << "u = " << u << " v = " << v << std::endl;

    iterCnt++;
    std::cout << "Iterator " << iterCnt << " times" << std::endl;
    // std::cout << "Levenberg-Marquardt parameter: " << std::endl << params_lm << std::endl << std::endl << std::endl;
    // break;
  }

  // step: 3. 输出归一化结果
  std::cout << "Levenberg-Marquardt parameter: " << std::endl
            << params_lm / params_lm(0) << std::endl
            << std::endl
            << std::endl;
}

int main(int argc, char **argv)
{
  // step: 1. 生成带噪声的椭圆数据
  int  data_cnts = 200;
  bool add_noise = false; // 控制数据是否有噪声

  // gt:
  // a = 1
  // b = 0.988215
  // c = 0.00119476
  // d = -784.547
  // e = -534.193
  // f = 213541

  // ================= case1 =================
  // double A = 3, a = 1, B = 2, b = 2;
  // (x-a)^2 / A^2 + (y-b)^2 / B^2 - 1 = 0 转以下方程
  // ax^2 + by^2 + cxy + dx + ey + f = 0
  // 4x^2 + 9y^2 - 8x - 36y + 4 = 0
  // a = 4 / 4 = 1
  // b = 9 / 4 = 2.25
  // c = 0
  // d = -8 / 4 = -2
  // e = -36 / 4 = -9
  // f = 4 / 4 = 1

  // ================= case2 =================
  double A = 3, a = 1, B = 1, b = 3;
  // x^2 + 9y^2 - 2x - 54y + 73 = 0
  // a = 1 / 73 = 0.0137
  // b = 9 / 73 = 0.123287671
  // c = 0
  // d = -2 / 73 = −0.02739726
  // e = -54 / 73 = −0.739726027
  // f = 73 / 73 = 1

  std::vector<cv::Point2d> input; // 上下对称数据
  // for (size_t i = 0; i < data_cnts; ++i)
  // {
  //   double x = 2 * A * ((random() % 100) / 100.0) - A;
  //   double s = std::sqrt((1 - (x - a) * (x - a) / (A * A)) * B * B);
  //   if (std::isnan(s))
  //     continue;

  //   double noise1 = (random() % 1000) / 10000.0;
  //   double y1     = add_noise ? b + s + noise1 : b + s;
  //   input.push_back(cv::Point2d(x, y1));

  //   double noise2 = (random() % 1000) / 10000.0;
  //   double y2     = add_noise ? b - s + noise2 : b - s;
  //   input.push_back(cv::Point2d(x, y2));
  // }
  
  std::ifstream file_x;
  file_x.open("../x.txt");
  std::ifstream file_y;
  file_y.open("../y.txt");
  std::string line_x, line_y;
  while (getline(file_x, line_x))
  {
    getline(file_y, line_y);
    double x = std::stod(line_x);
    double y = std::stod(line_y);
    std::cout << x << " " << y << std::endl;
    input.push_back(cv::Point2d(x, y));
  }

  // step: 2. LM 拟合
  LM(input, 200);

  return 0;
}